const tf = require('@tensorflow/tfjs-node');
const fs  = require('fs')
const path = require('path')
const mobilenet = require('@tensorflow-models/mobilenet');
const image = require('@tensorflow-models/image-recognition');  
console.log('开始了哦', mobilenet)
console.log('开始了哦1', image)


const typeList = {
    daisy: {
        key: 1,
        value: '雏菊'
    },
    dandelion: {
        key: 2,
        value: '蒲公英'
    },
    rose: {
        key: 3,
        value: '玫瑰'
    },
    sunflower: {
        key: 4,
        value: '向日葵'
    },
    tulip: {
        key: 5,
        value: '郁金香'
    }
}

/**
 * 创建一个未训练的模型
 * @returns {tf.Sequential}
 */
const createModel = () => {
  // 网络每一层的一些参数
  const kernel_size = [3, 3]
  const pool_size= [2, 2]
  const first_filters = 32
  const dropout_conv = 0.3
  const dropout_dense = 0.3
  // 建立一个模型
  const mModel = tf.sequential() // 获取创建的模型的类别
  // 卷积层1
  mModel.add(tf.layers.conv2d({
    inputShape: [96, 96, 3],
    filters: first_filters,
    kernelSize: kernel_size,
    activation: 'relu'
  }))
  // 卷积层2
  mModel.add(tf.layers.conv2d({
    filters: first_filters,
    kernelSize: kernel_size,
    activation: 'relu'
  }))
  // 池化层
  mModel.add(tf.layers.maxPooling2d({ poolSize: pool_size }))
  mModel.add(tf.layers.dropout({ rate: dropout_conv }))
  mModel.add(tf.layers.flatten())
  // 全连接层1
  mModel.add(tf.layers.dense({
    units: 256,
    activation: 'relu'
  }))
  mModel.add(tf.layers.dropout({
    rate: dropout_dense
  }))
  // 全连接层2
  mModel.add(tf.layers.dense({
    units: 3, // 类别个数
    activation: 'softmax'
  }))
  // 编译模型
  mModel.compile({
    optimizer: tf.train.adam(0.0001), // 定义优化器
    loss: 'categoricalCrossentropy', // 定义损失函数
    metrics: ['accuracy'] // 评价指标  
  })
  // 查看模型
  mModel.summary()
  console.log('查看模型')
  return mModel
}

const imageToTensor = (paths) => {
  // concat方法是将多个张量合并成一个张量
  return tf.concat(paths.map(path => {
    // 使用tidy方法，对张量进行内存控制，不然你直接读直接内存炸
    try {
        return tf.tidy(() => {
          return tf.node
          .decodeImage(fs.readFileSync(path)) // 把图片转换成张量
          .resizeNearestNeighbor([96, 96]) // 把图片的宽高转换成96x96
          .toFloat() // 转换成浮点型的张量
          .div(tf.scalar(255.0)) // 将图片张量进行归一化（这个自行去理解）
          .expandDims()
        })
    } catch {
        console.log('path', path)
    }
  }))
}
const getLabels = (paths) => {
  const mLabels = paths.map((path, index) => {
    const mTemp = path.split('/')
    const mFilename = mTemp[mTemp.length - 2]
    return typeList[mFilename].key
    // return 1
  })
  // mLabels输出的结果是只有0和1的数组，例如：[0,0,0,1,1,1,0,1,0...]
  // 接下来我们将这个数组转换成张量，这一行自行翻源代码理解哈
  console.log('mLabels', mLabels)
  // return tf.oneHot(tf.tensor1d(mLabels, 'int32'), Object.keys(typeList).length).toFloat()
  // tf.tensor1d 将数组转换成张量
  // tf.oneHot 将张量转换成合适的格式
  return tf.oneHot(tf.tensor1d(mLabels, 'int32'), 3).toFloat()
}

fs.readdir('../image/flowers', (err, files) => {
    if (err) {
        console.log('读取文件错误', err)
        return
    }
    console.log('文件目录', files)
    const prePath = '../image/flowers/'
    const filesList = files.reduce((prev, dir) => {
        if (dir === 'test') return prev
        let fileList = fs.readdirSync(prePath + dir)
        fileList = fileList.map(file => {
            return `${prePath + dir}/${file}`
        })
        // console.log('fileList', fileList.length)
        prev.push(...fileList)
        return prev
    }, [])
    console.log('filesList', filesList)
    console.log('filesList', path.join(__dirname, '../model/' + new Date().getTime()))
    tf.util.shuffle(filesList)
    const mTarinTensor = imageToTensor(filesList)
    const mTarinLabel = getLabels(filesList)
    console.log('mTarinLabel---------------')
    mTarinLabel.print()
    console.log(mTarinLabel.shape)
    console.log('mTarinLabel---------------end')
    console.log('mTarinTensor', mTarinTensor, mTarinLabel)
    // return
    const mModel = createModel()
    // 训练模型，只需要传入我们做好的两个张量，一个是样本一个是标注
    mModel.fit(mTarinTensor, mTarinLabel, {
      epochs: 15, // 训练15次，也就是这8005张图片要被训练15次
      batchSize: 32, // 每批图片要训练多少张
      callbacks: {
        onEpochEnd: () => {
          // 每一次训练结束后执行的操作，保存模型
          // mModel.save(path.join(__dirname, '../model/' + new Date().getTime()))
          mModel.save('file://' + new Date().getTime())
        }
      }
    })
})
// const mt = tf.node.decodeImage(fs.readFileSync('2122401867_cd86c5f114_n.jpg'))
// mt.print()
